import biotite.interface.rdkit as rdkit_interface
import biotite.structure.info as info
import rdkit.Chem.AllChem as Chem
import rdkit.Chem.Draw as Draw

atoms = info.residue("C")
atoms = atoms[atoms.element != "H"]

mol = rdkit_interface.to_mol(atoms, explicit_hydrogen=True)
Chem.Compute2DCoords(mol)
Draw.MolToFile(mol, "explicit.png")

mol = rdkit_interface.to_mol(atoms, explicit_hydrogen=False)
Chem.Compute2DCoords(mol)
Draw.MolToFile(mol, "non_explicit.png")

import biotite.structure as struc
from biotite.interface.rdkit.mol import to_mol, from_mol
import rdkit.Chem.AllChem as Chem
from typing import *

try:
    # Settings for debugging & interactive tests
    from rdkit.Chem.Draw import IPythonConsole

    IPythonConsole.kekulizeStructures = False
    IPythonConsole.drawOptions.addAtomIndices = False
    IPythonConsole.ipython_3d = False
    IPythonConsole.ipython_useSVG = True
    IPythonConsole.drawOptions.addStereoAnnotation = True
    IPythonConsole.molSize = 300, 200
except ImportError:
    pass

tyr = struc.info.residue("TYR")
tyr.chain_id[:] = "A"
tyr.res_id[:] = 17
tyr_no_h = tyr[tyr.element != "H"]

def has_radicals(mol, return_indices: bool = False) -> Union[bool, Tuple[bool, List[int]]]:
    if mol is None:
        raise ValueError("Input molecule is None")
        
    radical_indices = []
    
    for atom in mol.GetAtoms():
        # Get number of radical electrons
        num_radical_electrons = atom.GetNumRadicalElectrons()
        if num_radical_electrons > 0:
            radical_indices.append(atom.GetIdx())
    
    has_any_radicals = len(radical_indices) > 0
    
    if return_indices:
        return has_any_radicals, radical_indices
    return has_any_radicals

def is_sanitized(mol) -> bool:
    if mol is None:
        return False
        
    try:
        # Create a copy to avoid modifying the input
        mol_copy = Chem.Mol(mol)
        
        # Try to sanitize with all checks enabled
        # catchErrors=True prevents exceptions from being raised
        result = Chem.SanitizeMol(mol_copy, catchErrors=True)
        
        # SANITIZE_NONE (0) means all checks passed
        return result == Chem.SanitizeFlags.SANITIZE_NONE
        
    except Exception:
        return False
    
rad_color = lambda has_rad: '\033[32m' if not has_rad else '\033[31m'  # Green if False, Red if True
san_color = lambda is_san: '\033[32m' if is_san else '\033[31m'      # Green if True, Red if False
reset_color = '\033[0m'

i = 0
for m in [tyr_no_h, tyr]:
    for explicit_hydrogen in [True, False]:
        mol = to_mol(m, explicit_hydrogen=explicit_hydrogen)
        Chem.Compute2DCoords(mol)
        print(f"------- {i}: has_H={'H' in m.element}, {explicit_hydrogen=} -------")
        
        # Color the output based on the conditions
        print("before sanitize")
        has_rad = has_radicals(mol)
        is_san = is_sanitized(mol)
        print(f"{rad_color(has_rad)}has_radicals(mol): {has_rad}{reset_color}")
        print(f"{san_color(is_san)}is_sanitized(mol): {is_san}{reset_color}")
        display(mol);

        Chem.SanitizeMol(mol)
        print("after sanitize")
        has_rad = has_radicals(mol)
        is_san = is_sanitized(mol)
        print(f"{rad_color(has_rad)}has_radicals(mol): {has_rad}{reset_color}")
        print(f"{san_color(is_san)}is_sanitized(mol): {is_san}{reset_color}")
        display(mol);

        print("converted back into biotite structure")
        print(repr(from_mol(mol, conformer_id=0)))

        i += 1